﻿using System;
using System.Collections.Generic;
using UnityEngine;

namespace Common
{

    public class Pool<T> 
    {
        private readonly Stack<T> m_Stack;
        private readonly Action<T> m_ResetAction;
        private readonly Func<T> m_CreateFunc;
        
        public int Count => m_Stack.Count;
        public Pool(Func<T> createFunc, Action<T> resetAction)
        {
            m_Stack = new Stack<T>();
            m_CreateFunc = createFunc;
            m_ResetAction = resetAction;
        }
        
        public Pool(int capacity, Func<T> createFunc, Action<T> resetAction)
        {
            m_Stack = new Stack<T>(capacity);
            m_CreateFunc = createFunc;
            m_ResetAction = resetAction;
        }

        private void Grow()
        {
            if (m_CreateFunc == null) throw new Exception("CreateFunc is null");
            m_Stack.Push(m_CreateFunc());
        }

        public T Get()
        {
            if (m_Stack.Count == 0)
            {
                Grow();
            }

            return m_Stack.Pop();
        }

        public void Return(T element)
        {
            m_ResetAction?.Invoke(element);
            m_Stack.Push(element);
        }
    }


    public class SubClassPool<T>
    {
        private static readonly Dictionary<Type, Pool<T>> s_SubPools = new Dictionary<Type, Pool<T>>();
        
        public static TSub Get<TSub>() where TSub : class, T, new()
        {
            var type = typeof(TSub);
            if (!s_SubPools.TryGetValue(type, out var pool))
            {
                pool = new Pool<T>(() => new TSub(), null);
                s_SubPools.Add(type, pool);
            }

            return pool.Get() as TSub;
        }
        
        public static void Return<TSub>(TSub state) where TSub : class, T
        {
            var type = typeof(TSub);
            if (!s_SubPools.TryGetValue(type, out var pool))
            {
#if UNITY_EDITOR
                Debug.LogError("Return to a pool that doesn't exist. Type: " + type.Name);
#endif
                return;
            }

            pool.Return(state);
        }
    }
}